import ast
import logging
from typing import Dict, Union

from black import format_str, FileMode

from 程序核心代码.ast节点.节点定义.基础定义.ast节点基础定义 import ast节点基础模型
from 程序核心代码.ast节点.节点定义.基础定义.节点异常定义 import 节点异常

logger = logging.getLogger(__name__)

class 代码生成器:
    """代码生成器类，负责从AST节点生成格式化的源代码"""
    
    def __init__(self):
        self.缩进级别 = 0
        self.缩进字符 = "    "  # 4个空格作为缩进
        self.换行符 = "\n"
        self.注释映射: Dict[int, str] = {}  # 行号到注释的映射
        
    def 生成代码(self, 节点: Union[ast.AST, ast节点基础模型],
                 保留注释: bool = True,
                 格式化代码: bool = True) -> str:
        """从AST节点生成源代码"""
        try:
            # 如果输入是基础节点，获取其AST节点
            if isinstance(节点, ast节点基础模型):
                ast节点 = 节点._ast节点
            else:
                ast节点 = 节点
                
            # 提取原始注释（如果需要）
            if 保留注释:
                self._提取注释(ast节点)
                
            # 生成源代码
            源代码 = ast.unparse(ast节点)
            
            # 插入注释
            if 保留注释:
                源代码 = self._插入注释(源代码)
                
            # 格式化代码
            if 格式化代码:
                try:
                    源代码 = format_str(源代码, mode=FileMode())
                except Exception as e:
                    logger.warning(f"代码格式化失败: {str(e)}")
                    
            return 源代码
            
        except Exception as e:
            logger.error(f"代码生成失败: {str(e)}")
            raise 节点异常(f"代码生成失败: {str(e)}") from e
            
    def _提取注释(self, 节点: ast.AST) -> None:
        """从AST节点中提取注释信息"""
        for 子节点 in ast.walk(节点):
            if not hasattr(子节点, 'lineno'):
                continue
                
            行号 = 子节点.lineno
            
            # 提取节点的文档字符串
            if isinstance(子节点, (ast.FunctionDef, ast.ClassDef, ast.Module)):
                文档字符串 = ast.get_docstring(子节点)
                if 文档字符串:
                    self.注释映射[行号] = f'"""{文档字符串}"""'
                    
    def _插入注释(self, 源代码: str) -> str:
        """将注释插入到源代码中"""
        if not self.注释映射:
            return 源代码
            
        代码行 = 源代码.splitlines()
        结果行 = []
        
        for i, 行 in enumerate(代码行, 1):
            if i in self.注释映射:
                结果行.append(self.注释映射[i])
            结果行.append(行)
            
        return self.换行符.join(结果行)
        
    def 设置缩进(self, 缩进字符: str = "    ") -> None:
        """设置缩进字符"""
        self.缩进字符 = 缩进字符
        
    def 设置换行符(self, 换行符: str = "\n") -> None:
        """设置换行符"""
        self.换行符 = 换行符
        
    def 重置(self) -> None:
        """重置生成器状态"""
        self.缩进级别 = 0
        self.注释映射.clear() 